import os
import sys
#受体名字
receptor = sys.argv[1]
#配体名字
ligand = sys.argv[2]
#是否block抗体（配体是抗体的话）
antibody_block = sys.argv[3]

dir = receptor.split(".pdb")[0] + ligand.split(".pdb")[0]
os.system("./mark_sur {} receptor.pdb".format(receptor))
os.system("./mark_sur {} ligand.pdb".format(ligand))

if (antibody_block == "y") | (antibody_block == "Y"):
	name_out = "receptor_block.pdb"
	list_hv = []
	list_lv = []
	# block FR 区域
	for i in range(98):
		if (i + 1 < 31) | ((i + 1 > 35) & (i + 1 < 50)) | ((i + 1 > 66) & (i + 1 < 99)):
			list_hv.append(str(i + 1))
	for i in range(98):
		if (i + 1 < 24) | ((i + 1 > 34) & (i + 1 < 50)) | ((i + 1 > 56) & (i + 1 < 89)):
			list_lv.append(str(i + 1))
	# 读取 抗体pdb文件

	pdb = open("receptor.pdb", "r")
	pdb_line = [i for i in pdb]
	pdb_line_atom = []
	for i in pdb_line:
		if i[0:4] == 'ATOM':
			pdb_line_atom.append(i)
	list_hv.append(pdb_line_atom[0][20:23].strip())
	list_lv.append(pdb_line_atom[-1][20:23].strip())
	for i in range(len(pdb_line_atom)):
		if (pdb_line_atom[i][20:23].strip() in list_hv) & (pdb_line_atom[i][23:26].strip() in list_hv):
			pdb_line_atom[i] = pdb_line_atom[i][0:55] + "19 " + pdb_line_atom[i][58:-1] + pdb_line_atom[i][-1]
		elif (pdb_line_atom[i][20:23].strip() in list_lv) & (pdb_line_atom[i][23:26].strip() in list_lv):
			pdb_line_atom[i] = pdb_line_atom[i][0:55] + "19 " + pdb_line_atom[i][58:-1] + pdb_line_atom[i][-1]
	# print(pdb_line_atom[i][54:58].strip())
	# print(pdb_line_atom[i][20:23].strip())
	pdb_w = open(name_out, "w")
	pdb_w.writelines(pdb_line_atom)
	pdb_w.close()
	pdb.close()
	print("受体抗体蛋白 已封闭")
	print("zdock 开始对接")
	os.system("./zdock -R receptor_block.pdb -L ligand.pdb -o zdock.out -D 6 -N 2000")
else:
	print("zdock 开始对接")
	os.system("./zdock -R receptor.pdb -L ligand.pdb -o zdock.out -D 6 -N 2000")
os.system("mkdir {}".format(dir))
os.system("cp zdock.out ./{}".format(dir))
os.system("./create.pl ./{}/zdock.out".format(dir))



